from funasr import AutoModel
import jieba

# 初始化模型
model = AutoModel(model="paraformer-zh")

# 执行 ASR 识别
res = model.generate(input="../file/asr_example_zh.wav")
print(res)

def simple_split_sentences(text, max_len=10):
    """
    使用 jieba 对无标点中文文本进行断句，并限制每句话不超过 max_len 字
    """
    words = list(jieba.cut(text))
    sentences = []
    current_sentence = ""

    for word in words:
        if len(current_sentence) + len(word) <= max_len:
            current_sentence += word
        else:
            sentences.append(current_sentence)
            current_sentence = word

    if current_sentence:
        sentences.append(current_sentence)

    return sentences


def ms_to_srt_time(ms):
    """将毫秒转换为 SRT 时间格式"""
    seconds, ms = divmod(ms, 1000)
    minutes, seconds = divmod(seconds, 60)
    hours, minutes = divmod(minutes, 60)
    return f"{hours:02d}:{minutes:02d}:{seconds:02d},{ms:03d}"


def generate_srt(data):
    text = data[0]['text']
    timestamps = data[0]['timestamp']

    # 清理原始文本，去除空格
    cleaned_text = text.replace(' ', '')

    # 使用 jieba 断句
    sentences = simple_split_sentences(cleaned_text, max_len=12)

    srt_blocks = []
    ts_index = 0

    for i, sentence in enumerate(sentences):
        start_time = timestamps[ts_index][0]
        end_time = timestamps[-1][1] if i == len(sentences) - 1 else \
        timestamps[min(ts_index + len(sentence), len(timestamps) - 1) -1][1]

        srt_block = (
            f"{i + 1}\n"
            f"{ms_to_srt_time(start_time)} --> {ms_to_srt_time(end_time)}\n"
            f"{sentence}\n"
        )
        srt_blocks.append(srt_block)
        ts_index += max(1, len(sentence))

    return "\n".join(srt_blocks)


# 示例调用并输出到文件
srt_content = generate_srt(res)

with open("output.srt", "w", encoding="utf-8") as f:
    f.write(srt_content)

print("✅ SRT 文件已生成")
